import redis
import json
import os
from dotenv import load_dotenv
from langchain_core.messages import BaseMessage, message_to_dict, messages_from_dict
from config.config import Config

load_dotenv()

class MemoryStore:
    def __init__(self, redis_url: str = None):
        self.client = redis.Redis(
            host=Config.redis["REDIS_HOST"],
            port=int(Config.redis["REDIS_PORT"]),
            password=Config.redis["REDIS_PASSWORD"],
            db=int(Config.redis["REDIS_DB"]),
            decode_responses=False,  # 保持二进制格式
            socket_timeout=5,
            socket_connect_timeout=5,
            health_check_interval=30
        )
        # self.client = redis.Redis.from_url(
        #     redis_url or os.getenv("REDIS_URL"), 
        #     decode_responses=False
        # )
    
    def save_conversation(self, session_id: str, messages: list[BaseMessage]):
        """保存对话历史到Redis"""
        serialized = json.dumps([message_to_dict(m) for m in messages])
        # 设置30分钟过期时间
        self.client.setex(session_id, 1800, serialized)
    
    def load_conversation(self, session_id: str) -> list[BaseMessage]:
        """从Redis加载对话历史"""
        data = self.client.get(session_id)
        if not data:
            return []
        
        return messages_from_dict(json.loads(data))
    
    def add_message(self, session_id: str, message: BaseMessage):
        """添加单条消息到对话历史"""
        history = self.load_conversation(session_id)
        history.append(message)
        self.save_conversation(session_id, history)
        return history